#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebDriver管理工具
自动下载和管理与本地Chrome浏览器版本匹配的ChromeDriver
支持自动下载便携版Chromium，无需用户安装Chrome浏览器
"""

import os
import sys
import logging
import zipfile
import shutil
import urllib.request
from pathlib import Path
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from webdriver_manager.chrome import ChromeDriverManager


class ChromiumDownloader:
    """便携版Chromium自动下载器"""
    
    # Chromium下载地址（Chrome for Testing）
    CHROMIUM_DOWNLOAD_URLS = {
        'win64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/win64/chrome-win64.zip',
        'win32': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/win32/chrome-win32.zip',
        'mac64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/mac-x64/chrome-mac-x64.zip',
        'mac-arm64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/mac-arm64/chrome-mac-arm64.zip',
        'linux64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/linux64/chrome-linux64.zip',
    }
    
    # ChromeDriver下载地址
    CHROMEDRIVER_DOWNLOAD_URLS = {
        'win64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/win64/chromedriver-win64.zip',
        'win32': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/win32/chromedriver-win32.zip',
        'mac64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/mac-x64/chromedriver-mac-x64.zip',
        'mac-arm64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/mac-arm64/chromedriver-mac-arm64.zip',
        'linux64': 'https://storage.googleapis.com/chrome-for-testing-public/131.0.6778.87/linux64/chromedriver-linux64.zip',
    }
    
    def __init__(self, logger=None):
        self.logger = logger or self._create_default_logger()
        self.chromium_dir = Path.home() / '.shopspider' / 'chromium'
        self.chromium_dir.mkdir(parents=True, exist_ok=True)
    
    def _create_default_logger(self):
        """创建默认logger"""
        logger = logging.getLogger('ChromiumDownloader')
        logger.setLevel(logging.INFO)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        return logger
    
    def get_platform(self):
        """检测操作系统平台"""
        import platform
        
        if sys.platform.startswith('win'):
            # Windows：检测是32位还是64位
            if platform.machine().endswith('64'):
                return 'win64'
            else:
                return 'win32'
        elif sys.platform == 'darwin':
            # macOS：检测是Intel还是Apple Silicon
            machine = platform.machine().lower()
            if machine == 'arm64' or 'arm' in machine:
                return 'mac-arm64'  # Apple Silicon (M1/M2/M3)
            else:
                return 'mac64'      # Intel Mac
        elif sys.platform.startswith('linux'):
            # Linux：目前只支持64位
            return 'linux64'
        else:
            # 其他平台暂不支持
            return None
    
    def is_chromium_downloaded(self):
        """检查是否已下载Chromium"""
        platform = self.get_platform()
        if not platform:
            return False
        
        # 检查Chrome可执行文件是否存在
        chrome_path = self._get_chrome_executable_path(platform)
        return chrome_path and chrome_path.exists()
    
    def _get_chrome_executable_path(self, platform):
        """获取Chrome可执行文件的完整路径"""
        if platform == 'win64':
            return self.chromium_dir / 'chrome-win64' / 'chrome.exe'
        elif platform == 'win32':
            return self.chromium_dir / 'chrome-win32' / 'chrome.exe'
        elif platform == 'mac64':
            return self.chromium_dir / 'chrome-mac-x64' / 'Google Chrome for Testing.app' / 'Contents' / 'MacOS' / 'Google Chrome for Testing'
        elif platform == 'mac-arm64':
            return self.chromium_dir / 'chrome-mac-arm64' / 'Google Chrome for Testing.app' / 'Contents' / 'MacOS' / 'Google Chrome for Testing'
        elif platform == 'linux64':
            return self.chromium_dir / 'chrome-linux64' / 'chrome'
        else:
            return None
    
    def get_chromium_path(self):
        """获取Chromium可执行文件路径"""
        platform = self.get_platform()
        if not platform:
            return None
        
        return self._get_chrome_executable_path(platform)
    
    def get_chromedriver_path(self):
        """获取ChromeDriver可执行文件路径"""
        platform = self.get_platform()
        if not platform:
            return None
        
        if platform == 'win64':
            return self.chromium_dir / 'chromedriver-win64' / 'chromedriver.exe'
        elif platform == 'win32':
            return self.chromium_dir / 'chromedriver-win32' / 'chromedriver.exe'
        elif platform == 'mac64':
            return self.chromium_dir / 'chromedriver-mac-x64' / 'chromedriver'
        elif platform == 'mac-arm64':
            return self.chromium_dir / 'chromedriver-mac-arm64' / 'chromedriver'
        elif platform == 'linux64':
            return self.chromium_dir / 'chromedriver-linux64' / 'chromedriver'
        else:
            return None
    
    def download_chromium(self, progress_callback=None):
        """
        下载便携版Chromium
        
        Args:
            progress_callback: 进度回调函数 callback(current, total, percent)
            
        Returns:
            bool: 下载是否成功
        """
        platform = self.get_platform()
        if not platform:
            self.logger.error("❌ 不支持的操作系统平台")
            return False
        
        # 如果已下载，跳过
        if self.is_chromium_downloaded():
            self.logger.info("✅ Chromium已存在，跳过下载")
            return True
        
        try:
            self.logger.info("🌐 开始下载便携版Chromium...")
            self.logger.info(f"📂 目标目录: {self.chromium_dir}")
            
            # 下载Chrome
            chrome_url = self.CHROMIUM_DOWNLOAD_URLS[platform]
            chrome_zip = self.chromium_dir / 'chrome.zip'
            
            self.logger.info(f"⬇️  正在下载Chrome浏览器 (约120MB)...")
            self._download_with_progress(chrome_url, chrome_zip, progress_callback)
            
            self.logger.info("📦 正在解压Chrome...")
            with zipfile.ZipFile(chrome_zip, 'r') as zip_ref:
                zip_ref.extractall(self.chromium_dir)
            chrome_zip.unlink()  # 删除zip文件
            
            # 下载ChromeDriver
            driver_url = self.CHROMEDRIVER_DOWNLOAD_URLS[platform]
            driver_zip = self.chromium_dir / 'chromedriver.zip'
            
            self.logger.info("⬇️  正在下载ChromeDriver...")
            self._download_with_progress(driver_url, driver_zip)
            
            self.logger.info("📦 正在解压ChromeDriver...")
            with zipfile.ZipFile(driver_zip, 'r') as zip_ref:
                zip_ref.extractall(self.chromium_dir)
            driver_zip.unlink()  # 删除zip文件
            
            # Mac/Linux需要设置可执行权限
            if platform.startswith('mac') or platform.startswith('linux'):
                self._set_executable_permissions(platform)
            
            self.logger.info("✅ Chromium下载完成！")
            return True
            
        except Exception as e:
            self.logger.error(f"❌ 下载Chromium失败: {e}")
            return False
    
    def _set_executable_permissions(self, platform):
        """为Mac/Linux系统设置可执行权限"""
        try:
            import stat
            
            # 设置Chrome可执行权限
            chrome_path = self._get_chrome_executable_path(platform)
            if chrome_path and chrome_path.exists():
                chrome_path.chmod(chrome_path.stat().st_mode | stat.S_IEXEC)
                self.logger.info(f"✅ 设置Chrome可执行权限: {chrome_path}")
            
            # 设置ChromeDriver可执行权限
            driver_path = self.get_chromedriver_path()
            if driver_path and driver_path.exists():
                driver_path.chmod(driver_path.stat().st_mode | stat.S_IEXEC)
                self.logger.info(f"✅ 设置ChromeDriver可执行权限: {driver_path}")
                
        except Exception as e:
            self.logger.warning(f"⚠️ 设置可执行权限失败: {e}")
    
    def _download_with_progress(self, url, dest_path, progress_callback=None):
        """
        下载文件并显示进度
        
        Args:
            url: 下载URL
            dest_path: 保存路径
            progress_callback: 进度回调函数
        """
        def reporthook(count, block_size, total_size):
            if total_size > 0:
                percent = min(int(count * block_size * 100 / total_size), 100)
                current = min(count * block_size, total_size)
                
                # 格式化显示
                current_mb = current / 1024 / 1024
                total_mb = total_size / 1024 / 1024
                
                # 显示进度条
                bar_length = 30
                filled = int(bar_length * percent / 100)
                bar = '█' * filled + '░' * (bar_length - filled)
                
                print(f'\r  进度: [{bar}] {percent}% ({current_mb:.1f}/{total_mb:.1f} MB)', end='', flush=True)
                
                # 调用回调函数
                if progress_callback:
                    progress_callback(current, total_size, percent)
        
        try:
            urllib.request.urlretrieve(url, dest_path, reporthook=reporthook)
            print()  # 换行
        except Exception as e:
            print()  # 换行
            raise e


class WebDriverManager:
    """WebDriver管理器 - 自动管理ChromeDriver版本"""
    
    def __init__(self, logger=None, auto_download_chromium=True):
        """
        初始化WebDriver管理器
        
        Args:
            logger: 日志记录器，如果为None则创建默认logger
            auto_download_chromium: 是否自动下载Chromium（如果本地没有Chrome）
        """
        self.logger = logger or self._create_default_logger()
        self.auto_download_chromium = auto_download_chromium
        self.chromium_downloader = ChromiumDownloader(logger)
    
    def _create_default_logger(self):
        """创建默认的logger"""
        logger = logging.getLogger('WebDriverManager')
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def get_chrome_driver(self, options=None, headless=False):
        """
        获取Chrome WebDriver实例，自动匹配Chrome版本
        支持自动下载便携版Chromium（如果本地没有Chrome）
        
        Args:
            options: Chrome选项对象，如果为None则创建默认选项
            headless: 是否使用无头模式
            
        Returns:
            webdriver.Chrome: Chrome WebDriver实例
        """
        try:
            # 如果没有提供options，创建默认options
            if options is None:
                options = self.get_default_chrome_options(headless)
            
            # 策略1：尝试使用本地Chrome浏览器
            self.logger.info("🔍 检测本地Chrome浏览器...")
            chrome_version = self.get_chrome_version()
            
            if chrome_version:
                self.logger.info(f"✅ 检测到Chrome版本: {chrome_version}")
                
                try:
                    # 使用webdriver-manager自动管理ChromeDriver
                    service = Service(ChromeDriverManager().install())
                    driver = webdriver.Chrome(service=service, options=options)
                    self.logger.info("✅ 成功使用本地Chrome浏览器")
                    return driver
                except Exception as e:
                    self.logger.warning(f"⚠️ 使用本地Chrome失败: {e}")
            else:
                self.logger.info("⚠️ 未检测到本地Chrome浏览器")
            
            # 策略2：使用便携版Chromium
            if self.auto_download_chromium:
                self.logger.info("🌐 尝试使用便携版Chromium...")
                
                # 检查是否已下载
                if not self.chromium_downloader.is_chromium_downloaded():
                    self.logger.info("📥 首次运行需要下载Chromium (约120MB)...")
                    self.logger.info("💡 下载后会缓存，下次启动将秒开")
                    
                    # 下载Chromium
                    if not self.chromium_downloader.download_chromium():
                        raise Exception("Chromium下载失败")
                
                # 配置使用下载的Chromium
                chromium_path = self.chromium_downloader.get_chromium_path()
                chromedriver_path = self.chromium_downloader.get_chromedriver_path()
                
                if chromium_path and chromium_path.exists():
                    options.binary_location = str(chromium_path)
                    service = Service(str(chromedriver_path))
                    driver = webdriver.Chrome(service=service, options=options)
                    self.logger.info("✅ 成功使用便携版Chromium")
                    return driver
            
            # 策略3：尝试本地chromedriver
            self.logger.info("🔄 尝试使用本地ChromeDriver...")
            local_drivers = self._find_local_chromedriver()
            
            if local_drivers:
                for driver_path in local_drivers:
                    try:
                        self.logger.info(f"📂 尝试使用: {driver_path}")
                        service = Service(driver_path)
                        driver = webdriver.Chrome(service=service, options=options)
                        self.logger.info(f"✅ 成功使用本地ChromeDriver")
                        return driver
                    except Exception as local_e:
                        self.logger.warning(f"⚠️ 失败: {local_e}")
                        continue
            
            # 所有策略都失败
            raise Exception("无法初始化Chrome浏览器")
                
        except Exception as e:
            self.logger.error(f"❌ 无法初始化ChromeDriver: {e}")
            self.logger.error("💡 建议:")
            self.logger.error("   1. 检查网络连接（首次使用需要下载）")
            self.logger.error("   2. 手动安装Chrome浏览器")
            self.logger.error("   3. 运行: pip install webdriver-manager --upgrade")
            raise
    
    def get_default_chrome_options(self, headless=False):
        """
        获取默认的Chrome选项配置
        
        Args:
            headless: 是否使用无头模式
            
        Returns:
            Options: Chrome选项对象
        """
        chrome_options = Options()
        
        # 基础配置
        chrome_options.add_argument('--disable-blink-features=AutomationControlled')
        chrome_options.add_experimental_option('excludeSwitches', ['enable-automation'])
        chrome_options.add_experimental_option('useAutomationExtension', False)
        
        # 禁用GPU加速（提高兼容性）
        chrome_options.add_argument('--disable-gpu')
        
        # 禁用开发者模式警告
        chrome_options.add_experimental_option('excludeSwitches', ['enable-logging'])
        
        # 设置窗口大小
        chrome_options.add_argument('--window-size=1920,1080')
        
        # 无头模式
        if headless:
            chrome_options.add_argument('--headless')
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
        
        # 禁用图片加载（可选，提高速度）
        # prefs = {'profile.managed_default_content_settings.images': 2}
        # chrome_options.add_experimental_option('prefs', prefs)
        
        return chrome_options
    
    def _find_local_chromedriver(self):
        """
        查找项目中的本地chromedriver文件
        
        Returns:
            list: 找到的chromedriver路径列表
        """
        drivers = []
        
        # 根据平台确定可执行文件名
        if sys.platform.startswith('win'):
            driver_filename = 'chromedriver.exe'
        else:
            driver_filename = 'chromedriver'  # Mac/Linux
        
        # 搜索的目录列表
        search_dirs = [
            os.getcwd(),  # 当前目录
            os.path.dirname(__file__),  # 脚本所在目录
            os.path.join(os.getcwd(), 'PDD'),
            os.path.join(os.getcwd(), 'JD'),
            os.path.join(os.getcwd(), 'TaoBao'),
            os.path.join(os.getcwd(), 'DY'),
        ]
        
        # 搜索chromedriver
        for search_dir in search_dirs:
            if not os.path.exists(search_dir):
                continue
                
            driver_path = os.path.join(search_dir, driver_filename)
            if os.path.exists(driver_path) and driver_path not in drivers:
                drivers.append(driver_path)
        
        return drivers
    
    @staticmethod
    def get_chrome_version():
        """
        获取本地安装的Chrome浏览器版本
        
        Returns:
            str: Chrome版本号，如果获取失败返回None
        """
        try:
            import subprocess
            import re
            
            # Windows系统
            if os.name == 'nt':
                # 尝试从注册表获取
                try:
                    import winreg
                    key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, 
                                        r'Software\Google\Chrome\BLBeacon')
                    version, _ = winreg.QueryValueEx(key, 'version')
                    return version
                except:
                    pass
                
                # 尝试从可执行文件获取
                chrome_paths = [
                    r'C:\Program Files\Google\Chrome\Application\chrome.exe',
                    r'C:\Program Files (x86)\Google\Chrome\Application\chrome.exe',
                ]
                
                for chrome_path in chrome_paths:
                    if os.path.exists(chrome_path):
                        result = subprocess.run([chrome_path, '--version'], 
                                              capture_output=True, text=True)
                        match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
                        if match:
                            return match.group(1)
            
            # macOS系统
            elif os.name == 'posix' and sys.platform == 'darwin':
                chrome_paths = [
                    '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome',
                    '/Applications/Chrome.app/Contents/MacOS/Google Chrome',
                ]
                
                for chrome_path in chrome_paths:
                    if os.path.exists(chrome_path):
                        result = subprocess.run([chrome_path, '--version'], 
                                              capture_output=True, text=True)
                        match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
                        if match:
                            return match.group(1)
            
            # Linux系统
            elif os.name == 'posix' and sys.platform.startswith('linux'):
                chrome_commands = ['google-chrome', 'google-chrome-stable', 'chromium-browser']
                
                for cmd in chrome_commands:
                    try:
                        result = subprocess.run([cmd, '--version'], 
                                              capture_output=True, text=True)
                        match = re.search(r'(\d+\.\d+\.\d+\.\d+)', result.stdout)
                        if match:
                            return match.group(1)
                    except FileNotFoundError:
                        continue
            
        except Exception as e:
            logging.warning(f"无法获取Chrome版本: {e}")
        
        return None


# 便捷函数
def get_chrome_driver(headless=False, logger=None, custom_options=None, auto_download_chromium=True):
    """
    便捷函数：快速获取Chrome WebDriver
    自动处理Chrome/Chromium检测和下载
    
    Args:
        headless: 是否使用无头模式
        logger: 日志记录器
        custom_options: 自定义Chrome选项
        auto_download_chromium: 是否自动下载Chromium（如果本地没有Chrome）
        
    Returns:
        webdriver.Chrome: Chrome WebDriver实例
    """
    manager = WebDriverManager(logger, auto_download_chromium=auto_download_chromium)
    return manager.get_chrome_driver(options=custom_options, headless=headless)


if __name__ == "__main__":
    # 测试代码
    print("🧪 测试WebDriver管理器（支持跨平台自动下载Chromium）")  # 支持跨平台
    print("=" * 60)
    
    # 检测平台
    downloader = ChromiumDownloader()
    platform = downloader.get_platform()
    platform_names = {
        'win64': 'Windows 64位',
        'win32': 'Windows 32位', 
        'mac64': 'macOS Intel',
        'mac-arm64': 'macOS Apple Silicon (M1/M2/M3)',
        'linux64': 'Linux 64位'
    }
    
    print(f"\n🖥️  当前平台: {platform_names.get(platform, f'未知平台 ({platform})')}")
    
    if not platform:
        print("❌ 不支持的平台，程序退出")
        exit(1)
    
    # 检测Chrome版本
    print("\n1️⃣ 检测本地Chrome浏览器...")
    chrome_version = WebDriverManager.get_chrome_version()
    if chrome_version:
        print(f"   ✅ 检测到Chrome版本: {chrome_version}")
    else:
        print("   ⚠️ 未检测到本地Chrome浏览器")
        print("   💡 将自动下载便携版Chromium")
    
    # 检测Chromium
    print("\n2️⃣ 检测便携版Chromium...")
    if downloader.is_chromium_downloaded():
        print(f"   ✅ Chromium已下载")
        print(f"   📂 位置: {downloader.get_chromium_path()}")
    else:
        print("   ⚠️ Chromium未下载（首次运行会自动下载）")
        if platform.startswith('mac') or platform.startswith('linux'):
            print("   💡 Mac/Linux版本约150MB，Windows约120MB")
    
    # 测试获取WebDriver
    print("\n3️⃣ 测试获取WebDriver...")
    print("   （如果没有Chrome，将自动下载Chromium，请稍候...）")
    print()
    
    try:
        driver = get_chrome_driver(headless=False, auto_download_chromium=True)
        print("\n   ✅ WebDriver初始化成功")
        
        print("   🌐 测试访问百度...")
        driver.get("https://www.baidu.com")
        print(f"   ✅ 成功访问网页，标题: {driver.title}")
        
        import time
        time.sleep(2)
        
        driver.quit()
        print("   ✅ WebDriver已关闭")
        print("\n" + "=" * 60)
        print("✅ 所有测试通过！")
        print(f"✅ 当前平台 ({platform_names.get(platform)}) 完全支持！")
        
    except Exception as e:
        print(f"\n   ❌ 测试失败: {e}")
        import traceback
        traceback.print_exc() 